{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# {class}`~torch.nn.PixelShuffle`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch.autograd.grad_mode.set_grad_enabled at 0x7f6813e208c0>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from tvm import relay\n",
    "import tvm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "PyTorch has 1 inputs and input_infos lists 4.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m input_shapes \u001b[38;5;241m=\u001b[39m [(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mx_\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mindex\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, shape) \u001b[38;5;28;01mfor\u001b[39;00m index, shape \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(input_shape)]\n\u001b[1;32m      5\u001b[0m trace \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mjit\u001b[38;5;241m.\u001b[39mtrace(model, [torch\u001b[38;5;241m.\u001b[39mrand(input_shape)\u001b[38;5;241m.\u001b[39mfloat()])\n\u001b[0;32m----> 6\u001b[0m mod, params \u001b[38;5;241m=\u001b[39m relay\u001b[38;5;241m.\u001b[39mfrontend\u001b[38;5;241m.\u001b[39mfrom_pytorch(\n\u001b[1;32m      7\u001b[0m     trace,\n\u001b[1;32m      8\u001b[0m     input_shapes,\n\u001b[1;32m      9\u001b[0m )\n\u001b[1;32m     10\u001b[0m target \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mllvm\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     11\u001b[0m dev \u001b[38;5;241m=\u001b[39m tvm\u001b[38;5;241m.\u001b[39mdevice(target, \u001b[38;5;241m0\u001b[39m)\n",
      "File \u001b[0;32m/media/pc/data/board/arria10/lxw/tasks/tvm-ai/python/tvm/relay/frontend/pytorch.py:5391\u001b[0m, in \u001b[0;36mfrom_pytorch\u001b[0;34m(script_module, input_infos, custom_convert_map, default_dtype, use_parser_friendly_name, keep_quantized_weight, export_renamed_c_graph_path, preserve_pytorch_scopes)\u001b[0m\n\u001b[1;32m   5389\u001b[0m is_module \u001b[38;5;241m=\u001b[39m \u001b[38;5;28misinstance\u001b[39m(script_module, torch\u001b[38;5;241m.\u001b[39mjit\u001b[38;5;241m.\u001b[39mScriptModule)\n\u001b[1;32m   5390\u001b[0m params \u001b[38;5;241m=\u001b[39m script_module\u001b[38;5;241m.\u001b[39mstate_dict() \u001b[38;5;28;01mif\u001b[39;00m is_module \u001b[38;5;28;01melse\u001b[39;00m {}\n\u001b[0;32m-> 5391\u001b[0m outputs \u001b[38;5;241m=\u001b[39m _get_relay_input_vars(\n\u001b[1;32m   5392\u001b[0m     graph, input_infos, prelude, default_dtype\u001b[38;5;241m=\u001b[39mdefault_dtype, is_module\u001b[38;5;241m=\u001b[39mis_module\n\u001b[1;32m   5393\u001b[0m )\n\u001b[1;32m   5395\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m use_parser_friendly_name:\n\u001b[1;32m   5396\u001b[0m     new_names \u001b[38;5;241m=\u001b[39m [key\u001b[38;5;241m.\u001b[39mreplace(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mfor\u001b[39;00m key \u001b[38;5;129;01min\u001b[39;00m params\u001b[38;5;241m.\u001b[39mkeys()]\n",
      "File \u001b[0;32m/media/pc/data/board/arria10/lxw/tasks/tvm-ai/python/tvm/relay/frontend/pytorch.py:5069\u001b[0m, in \u001b[0;36m_get_relay_input_vars\u001b[0;34m(graph, input_infos, prelude, is_module, default_dtype)\u001b[0m\n\u001b[1;32m   5067\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(graph_inputs) \u001b[38;5;241m!=\u001b[39m \u001b[38;5;28mlen\u001b[39m(input_infos):\n\u001b[1;32m   5068\u001b[0m     msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPyTorch has \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(graph_inputs)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m inputs and input_infos lists \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(input_infos)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m-> 5069\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(msg)\n\u001b[1;32m   5071\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mget_relay_ty\u001b[39m(ishape, itype, pt_type):\n\u001b[1;32m   5072\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m pt_type\u001b[38;5;241m.\u001b[39mkind() \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTensorType\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
      "\u001b[0;31mRuntimeError\u001b[0m: PyTorch has 1 inputs and input_infos lists 4."
     ]
    }
   ],
   "source": [
    "input_shape = [1, 144, 16, 16]\n",
    "torch.set_grad_enabled(False)\n",
    "model = torch.nn.PixelShuffle(2).float().eval()\n",
    "input_shapes = [(f\"x_{index}\", shape) for index, shape in enumerate(input_shape)]\n",
    "trace = torch.jit.trace(model, [torch.rand(input_shape).float()])\n",
    "mod, params = relay.frontend.from_pytorch(\n",
    "    trace,\n",
    "    input_shapes,\n",
    ")\n",
    "target = \"llvm\"\n",
    "dev = tvm.device(target, 0)\n",
    "exe = relay.create_executor(\n",
    "    \"vm\", mod=mod, params=params, device=dev, target=target\n",
    ").evaluate()\n",
    "input_names = [f\"x_{idx}\" for idx, _ in enumerate(inputs)]\n",
    "compiled_input = dict(zip(input_names, [inp.clone().cpu().numpy() for inp in inputs]))\n",
    "result = exe(**compiled_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tvm.nd.NDArray shape=(1, 36, 32, 32), cpu(0)>\n",
       "array([[[[0.45167208, 0.4751528 , 0.66562855, ..., 0.10102624,\n",
       "          0.47107768, 0.98546124],\n",
       "         [0.50281024, 0.9156416 , 0.97454673, ..., 0.33834815,\n",
       "          0.99482936, 0.1813637 ],\n",
       "         [0.36587209, 0.77844733, 0.9402926 , ..., 0.07727742,\n",
       "          0.5027055 , 0.03256208],\n",
       "         ...,\n",
       "         [0.6315825 , 0.32724112, 0.1491878 , ..., 0.38751185,\n",
       "          0.71438605, 0.8736542 ],\n",
       "         [0.23701996, 0.93006486, 0.79464227, ..., 0.74860734,\n",
       "          0.10734552, 0.49250185],\n",
       "         [0.30654496, 0.90148944, 0.536291  , ..., 0.4301917 ,\n",
       "          0.1279993 , 0.517763  ]],\n",
       "\n",
       "        [[0.70649207, 0.01436931, 0.9452317 , ..., 0.50734496,\n",
       "          0.15042186, 0.75830054],\n",
       "         [0.81990063, 0.32498264, 0.7660303 , ..., 0.13331008,\n",
       "          0.12531197, 0.8993785 ],\n",
       "         [0.84151685, 0.15428996, 0.74632305, ..., 0.0347752 ,\n",
       "          0.35647523, 0.41170758],\n",
       "         ...,\n",
       "         [0.21683067, 0.25745237, 0.31181848, ..., 0.6034938 ,\n",
       "          0.735219  , 0.11763048],\n",
       "         [0.4092958 , 0.16630328, 0.9446607 , ..., 0.56538206,\n",
       "          0.6465735 , 0.93971133],\n",
       "         [0.06048411, 0.64390385, 0.5892658 , ..., 0.5724352 ,\n",
       "          0.5564407 , 0.21722358]],\n",
       "\n",
       "        [[0.03172588, 0.9610467 , 0.5827595 , ..., 0.34336174,\n",
       "          0.2407111 , 0.31925297],\n",
       "         [0.02371919, 0.8970315 , 0.26967138, ..., 0.13132441,\n",
       "          0.02313513, 0.24149656],\n",
       "         [0.88821566, 0.94423455, 0.24703354, ..., 0.11735535,\n",
       "          0.23171204, 0.77740663],\n",
       "         ...,\n",
       "         [0.23393357, 0.26179588, 0.5086286 , ..., 0.05173749,\n",
       "          0.9915401 , 0.98909074],\n",
       "         [0.21915352, 0.6296667 , 0.9010871 , ..., 0.9190174 ,\n",
       "          0.6990261 , 0.59724325],\n",
       "         [0.85011995, 0.27818072, 0.9096407 , ..., 0.06759602,\n",
       "          0.4988358 , 0.80241007]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[0.2965725 , 0.47470254, 0.87060493, ..., 0.8145422 ,\n",
       "          0.01968753, 0.41305685],\n",
       "         [0.3703339 , 0.16386402, 0.50278574, ..., 0.11774951,\n",
       "          0.90823394, 0.5985948 ],\n",
       "         [0.9500081 , 0.5035085 , 0.55892885, ..., 0.29899144,\n",
       "          0.22128135, 0.24032354],\n",
       "         ...,\n",
       "         [0.11933684, 0.14329827, 0.57943916, ..., 0.4280916 ,\n",
       "          0.54278034, 0.6509579 ],\n",
       "         [0.7960784 , 0.6684333 , 0.86588925, ..., 0.8933825 ,\n",
       "          0.69731855, 0.02392757],\n",
       "         [0.65791446, 0.9239144 , 0.7943546 , ..., 0.4539765 ,\n",
       "          0.19385767, 0.70213825]],\n",
       "\n",
       "        [[0.68747175, 0.8038203 , 0.66696763, ..., 0.6971842 ,\n",
       "          0.03665125, 0.9873148 ],\n",
       "         [0.94062805, 0.51540786, 0.63104075, ..., 0.10697812,\n",
       "          0.56769574, 0.28505236],\n",
       "         [0.04616427, 0.40854114, 0.40203577, ..., 0.31057447,\n",
       "          0.12055659, 0.3064711 ],\n",
       "         ...,\n",
       "         [0.35712874, 0.9227293 , 0.9166052 , ..., 0.65245557,\n",
       "          0.33464062, 0.93004066],\n",
       "         [0.7238047 , 0.88155776, 0.6993761 , ..., 0.71103656,\n",
       "          0.77516603, 0.37277114],\n",
       "         [0.95055205, 0.01364166, 0.8506952 , ..., 0.4274944 ,\n",
       "          0.17991948, 0.40492946]],\n",
       "\n",
       "        [[0.12583238, 0.39497036, 0.6986206 , ..., 0.8411366 ,\n",
       "          0.45737535, 0.06205308],\n",
       "         [0.7111632 , 0.43545884, 0.00555933, ..., 0.6724145 ,\n",
       "          0.2567408 , 0.6559393 ],\n",
       "         [0.3909461 , 0.9393658 , 0.63789135, ..., 0.6520382 ,\n",
       "          0.63449484, 0.21824509],\n",
       "         ...,\n",
       "         [0.8328597 , 0.58889294, 0.54670507, ..., 0.32606977,\n",
       "          0.25905436, 0.1458506 ],\n",
       "         [0.88871354, 0.3285848 , 0.71468323, ..., 0.52079886,\n",
       "          0.90121067, 0.9315291 ],\n",
       "         [0.07994634, 0.7397775 , 0.89741504, ..., 0.3094415 ,\n",
       "          0.5845787 , 0.17412347]]]], dtype=float32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
